{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOcMk1b3oNBRsuZ7B6aSsFY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sugarforever/LangChain-Tutorials/blob/main/langgraph_nodes_edges.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "id": "--5ZK7q9yi0l"
      },
      "outputs": [],
      "source": [
        "!pip install -q -U langchain langchain_openai langgraph google-search-results"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from google.colab import userdata"
      ],
      "metadata": {
        "id": "uwEe2bRl2OrV"
      },
      "execution_count": 82,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "os.environ['SERPAPI_API_KEY'] = userdata.get('GOOGLE_API_KEY')\n",
        "os.environ['OPENAI_API_KEY'] = userdata.get('OPENAI_API_KEY')\n",
        "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
        "os.environ[\"LANGCHAIN_PROJECT\"] = \"LangGraph\"\n",
        "os.environ[\"LANGCHAIN_API_KEY\"] = userdata.get('LANGSMITH_API_KEY')"
      ],
      "metadata": {
        "id": "TOLrjVQJ2Iyq"
      },
      "execution_count": 83,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_community.utilities import SerpAPIWrapper"
      ],
      "metadata": {
        "id": "7k4vCN42zAmb"
      },
      "execution_count": 84,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "search = SerpAPIWrapper()\n",
        "\n",
        "search.run(\"Obama's first name?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 36
        },
        "id": "pI2sLQxlzDOK",
        "outputId": "c8065154-fe4a-4602-dcf8-4213d0c04fbc"
      },
      "execution_count": 85,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'Barack Hussein Obama II'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 85
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import functools, operator, requests, os, json\n",
        "from langchain.agents import AgentExecutor, create_openai_tools_agent\n",
        "from langchain_core.messages import BaseMessage, HumanMessage\n",
        "from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser\n",
        "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
        "from langgraph.graph import StateGraph, END\n",
        "from langchain.tools import tool\n",
        "from langchain_openai import ChatOpenAI\n",
        "from typing import Annotated, Any, Dict, List, Optional, Sequence, TypedDict"
      ],
      "metadata": {
        "id": "tKCnieM72lBq"
      },
      "execution_count": 86,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "llm = ChatOpenAI(model=\"gpt-4-turbo-preview\")"
      ],
      "metadata": {
        "id": "KEyXHFis2fC2"
      },
      "execution_count": 87,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_core.messages import (\n",
        "    AIMessage,\n",
        "    BaseMessage,\n",
        "    ChatMessage,\n",
        "    FunctionMessage,\n",
        "    HumanMessage,\n",
        "    SystemMessage\n",
        ")\n",
        "\n",
        "@tool(\"web_search\")\n",
        "def web_search(query: str) -> str:\n",
        "    \"\"\"Search with Google SERP API by a query\"\"\"\n",
        "    search = SerpAPIWrapper()\n",
        "    return search.run(query)\n",
        "\n",
        "@tool(\"twitter_writer\")\n",
        "def write_tweet(content: str) -> str:\n",
        "    \"\"\"Based a piece of content, write a tweet.\"\"\"\n",
        "    chat = ChatOpenAI()\n",
        "    messages = [\n",
        "      SystemMessage(\n",
        "          content=\"You are a Twitter account operator.\"\n",
        "                  \" You are responsible for writing a tweet based on the content given.\"\n",
        "                  \" You should follow the Twitter policy and make sure each tweet has no more than 140 characters.\"\n",
        "      ),\n",
        "      HumanMessage(\n",
        "          content=content\n",
        "      ),\n",
        "    ]\n",
        "    response = chat(messages)\n",
        "    return response.content"
      ],
      "metadata": {
        "id": "fRUawUK622xO"
      },
      "execution_count": 88,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class AgentState(TypedDict):\n",
        "    # The annotation tells the graph that new messages will always\n",
        "    # be added to the current states\n",
        "    messages: Annotated[Sequence[BaseMessage], operator.add]\n",
        "    # The 'next' field indicates where to route to next\n",
        "    next: str"
      ],
      "metadata": {
        "id": "dX9Q74L8IK8I"
      },
      "execution_count": 90,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def create_agent(llm: ChatOpenAI, tools: list, system_prompt: str):\n",
        "    prompt = ChatPromptTemplate.from_messages(\n",
        "        [\n",
        "            (\n",
        "                \"system\",\n",
        "                system_prompt,\n",
        "            ),\n",
        "            MessagesPlaceholder(variable_name=\"messages\"),\n",
        "            MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
        "        ]\n",
        "    )\n",
        "    agent = create_openai_tools_agent(llm, tools, prompt)\n",
        "    executor = AgentExecutor(agent=agent, tools=tools)\n",
        "    return executor\n",
        "\n",
        "def agent_node(state, agent, name):\n",
        "    result = agent.invoke(state)\n",
        "    return {\"messages\": [HumanMessage(content=result[\"output\"], name=name)]}"
      ],
      "metadata": {
        "id": "t_qXl_Hz49KR"
      },
      "execution_count": 91,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "members = [\"Search_Engine\", \"Twitter_Writer\"]\n",
        "system_prompt = (\n",
        "    \"You are a supervisor tasked with managing a conversation between the\"\n",
        "    \" following workers:  {members}. Given the following user request,\"\n",
        "    \" respond with the worker to act next. Each worker will perform a\"\n",
        "    \" task and respond with their results and status. When finished,\"\n",
        "    \" respond with FINISH.\"\n",
        ")\n",
        "\n",
        "options = [\"FINISH\"] + members\n",
        "# Using openai function calling can make output parsing easier for us\n",
        "function_def = {\n",
        "    \"name\": \"route\",\n",
        "    \"description\": \"Select the next role.\",\n",
        "    \"parameters\": {\n",
        "        \"title\": \"routeSchema\",\n",
        "        \"type\": \"object\",\n",
        "        \"properties\": {\n",
        "            \"next\": {\n",
        "                \"title\": \"Next\",\n",
        "                \"anyOf\": [\n",
        "                    {\"enum\": options},\n",
        "                ],\n",
        "            }\n",
        "        },\n",
        "        \"required\": [\"next\"],\n",
        "    },\n",
        "}\n",
        "prompt = ChatPromptTemplate.from_messages(\n",
        "    [\n",
        "        (\"system\", system_prompt),\n",
        "        MessagesPlaceholder(variable_name=\"messages\"),\n",
        "        (\n",
        "            \"system\",\n",
        "            \"Given the conversation above, who should act next?\"\n",
        "            \" Or should we FINISH? Select one of: {options}\",\n",
        "        ),\n",
        "    ]\n",
        ").partial(options=str(options), members=\", \".join(members))\n",
        "\n",
        "supervisor_chain = (\n",
        "    prompt\n",
        "    | llm.bind_functions(functions=[function_def], function_call=\"route\")\n",
        "    | JsonOutputFunctionsParser()\n",
        ")"
      ],
      "metadata": {
        "id": "5KECklFd5z1J"
      },
      "execution_count": 92,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "search_engine_agent = create_agent(llm, [web_search], \"You are a web search engine.\")\n",
        "search_engine_node = functools.partial(agent_node, agent=search_engine_agent, name=\"Search_Engine\")\n",
        "\n",
        "twitter_operator_agent = create_agent(llm, [write_tweet], \"You are responsible for writing a tweet based on the content given.\")\n",
        "twitter_operator_node = functools.partial(agent_node, agent=twitter_operator_agent, name=\"Twitter_Writer\")\n",
        "\n",
        "workflow = StateGraph(AgentState)\n",
        "workflow.add_node(\"Search_Engine\", search_engine_node)\n",
        "workflow.add_node(\"Twitter_Writer\", twitter_operator_node)\n",
        "workflow.add_node(\"supervisor\", supervisor_chain)"
      ],
      "metadata": {
        "id": "vZID5tQv8E3n"
      },
      "execution_count": 93,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for member in members:\n",
        "    workflow.add_edge(member, \"supervisor\")\n",
        "\n",
        "conditional_map = {k: k for k in members}\n",
        "conditional_map[\"FINISH\"] = END\n",
        "workflow.add_conditional_edges(\"supervisor\", lambda x: x[\"next\"], conditional_map)\n",
        "\n",
        "workflow.set_entry_point(\"supervisor\")\n",
        "\n",
        "graph = workflow.compile()"
      ],
      "metadata": {
        "id": "LWBp9K739gt8"
      },
      "execution_count": 94,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for s in graph.stream(\n",
        "    {\n",
        "        \"messages\": [\n",
        "            HumanMessage(content=\"Write a tweet about LangChain news\")\n",
        "        ]\n",
        "    }\n",
        "):\n",
        "    if \"__end__\" not in s:\n",
        "        print(s)\n",
        "        print(\"----\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e9qi4pAX96Tf",
        "outputId": "f801546e-8406-48ab-cd35-bb567df5a5b2"
      },
      "execution_count": 95,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'supervisor': {'next': 'Search_Engine'}}\n",
            "----\n",
            "{'Search_Engine': {'messages': [HumanMessage(content=\"🚀 Exciting news from LangChain! 🌟\\n\\nWe just launched LangGraph, a revolutionary tool to customize your Agent Runtime, marking a significant milestone in our journey. Also, we're thrilled to announce the release of langchain 0.1.0, our first stable version that's fully backward compatible. 🎉\\n\\nStay tuned for more updates on how we're transforming the AI ecosystem. #LangChain #Innovation #AI\\n\\n[Week of 1/22/24]\", name='Search_Engine')]}}\n",
            "----\n",
            "{'supervisor': {'next': 'Twitter_Writer'}}\n",
            "----\n",
            "{'Twitter_Writer': {'messages': [HumanMessage(content='🚀 Exciting news from LangChain! 🌟 Introducing LangGraph, a customizable Agent Runtime tool, and langchain 0.1.0, our first stable release. #LangChain #Innovation #AI', name='Twitter_Writer')]}}\n",
            "----\n",
            "{'supervisor': {'next': 'FINISH'}}\n",
            "----\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "_SRgFulWAQ9R"
      },
      "execution_count": 95,
      "outputs": []
    }
  ]
}